# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
from typing import Union, List
import random
from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig

@GENERATOR_REGISTRY.register("ascend_generate_moe_token_permute_with_routing_map_grad")
class MoeTokenUnpermuteWithRoutingMapGradGenerator(CaseGenerator):
    def __init__(self, config):
        super().__init__(config)

    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        token_num = random.randint(16,1024*10)
        hidden_size = random.randint(1,1024*10)
        expert_num = random.randint(16,256*4)
        top_k = random.randint(1,16)
        capacity = random.randint(1,16)
        data_type = case_config.inputs[0].dtype
        case_config.inputs[1].dtype = data_type
        padded_mode_list = [True,False]
        padded_mode = random.choice(padded_mode_list)
        if padded_mode:
            case_config.inputs[0].shape = [expert_num * capacity, hidden_size]
            case_config.inputs[1].shape = [expert_num * capacity]
            case_config.inputs[2].shape = [expert_num * capacity]
        else:
            case_config.inputs[0].shape = [token_num * top_k, hidden_size]
            case_config.inputs[1].shape = [token_num * top_k]
            case_config.inputs[2].shape = [token_num * top_k]
        case_config.inputs[3].shape = [token_num, expert_num]
        case_config.inputs[4].range_values = expert_num
        case_config.inputs[5].range_values = token_num
        
        case_config.inputs[6].range_values = padded_mode
        
        return case_config